use crate::lexer::{Token, TokenKind};

// ------------------ AST start --------------------
pub struct Ast {
  statements: Vec<ASTStatement>,
}

impl Ast {
  pub fn new() -> Self {
    Self {
      statements: Vec::new(),
    }
  }

  pub fn add_statement(&mut self, statement: ASTStatement) {
    self.statements.push(statement)
  }

  pub fn visit(&mut self, visitor: &mut dyn ASTVisitor) {
    for statement in &self.statements {
      visitor.visit_statement(statement);
    }
  }

  pub fn visualize(&mut self) {
    let mut printer = ASTPrinter { indent: 0 };
    self.visit(&mut printer);
  }
}

pub trait ASTVisitor {
  fn do_visit_statement(&mut self, statement: &ASTStatement) {
    match &statement.kind {
      ASTStatementKind::Expression(expr) => {
        self.visit_expression(expr);
      }
    }
  }

  fn do_visit_expression(&mut self, expression: &ASTExpression) {
    match &expression.kind {
      ASTExpressionKind::Number(number) => self.visit_number(number),
      ASTExpressionKind::Binary(expr) => self.visit_binary_expression(expr),
      ASTExpressionKind::Parenthesized(expr) => self.visit_parenthesized_expression(expr),
      ASTExpressionKind::Unary(expr) => self.visit_unary_expression(expr),
      ASTExpressionKind::Boolean(expr) => self.visit_boolean_expression(expr),
    }
  }

  fn visit_statement(&mut self, statement: &ASTStatement) {
    self.do_visit_statement(statement);
  }

  fn visit_expression(&mut self, expression: &ASTExpression) {
    self.do_visit_expression(expression);
  }

  fn visit_number(&mut self, number: &ASTNumberExpression);

  fn visit_binary_expression(&mut self, binary_expression: &ASTBinaryExpression) {
    self.visit_expression(&binary_expression.lhs);
    self.visit_expression(&binary_expression.rhs);
  }

  fn visit_parenthesized_expression(
    &mut self,
    parenthesized_expression: &ASTParenthesizedExpression,
  ) {
    self.visit_expression(&parenthesized_expression.expression);
  }

  fn visit_unary_expression(&mut self, unary_expression: &ASTUnaryExpression);

  fn visit_boolean_expression(&mut self, boolean: &ASTBooleanExpression);
}

pub struct ASTPrinter {
  indent: usize,
}

const LEVEL_INDENT: usize = 2;

impl ASTVisitor for ASTPrinter {
  fn visit_number(&mut self, number: &ASTNumberExpression) {
    self.print_with_indent(&format!("Number: {:#?}", number.number));
  }

  fn visit_statement(&mut self, statement: &ASTStatement) {
    self.print_with_indent("Statement:");
    self.indent += LEVEL_INDENT;
    Self::do_visit_statement(self, statement);
    self.indent -= LEVEL_INDENT;
  }

  fn visit_expression(&mut self, expression: &ASTExpression) {
    self.print_with_indent("Expression:");
    self.indent += LEVEL_INDENT;
    Self::do_visit_expression(self, expression);
    self.indent -= LEVEL_INDENT;
  }

  fn visit_binary_expression(&mut self, binary_expression: &ASTBinaryExpression) {
    self.print_with_indent("Binary Expression:");
    self.indent += LEVEL_INDENT;
    self.print_with_indent(&format!("Operator: {:?}", binary_expression.operator.kind));
    self.visit_expression(&binary_expression.lhs);
    self.visit_expression(&binary_expression.rhs);
    self.indent -= LEVEL_INDENT;
  }

  fn visit_parenthesized_expression(
    &mut self,
    parenthesized_expression: &ASTParenthesizedExpression,
  ) {
    self.print_with_indent("Parenthesized Expression:");
    self.indent += LEVEL_INDENT;
    self.visit_expression(&parenthesized_expression.expression);
    self.indent -= LEVEL_INDENT;
  }

  fn visit_unary_expression(&mut self, unary_expression: &ASTUnaryExpression) {
    self.print_with_indent("Unary Expression:");
    self.indent += LEVEL_INDENT;
    self.print_with_indent(&format!("Operator: {:?}", unary_expression.operator.kind));
    self.visit_expression(&unary_expression.operand);
    self.indent -= LEVEL_INDENT;
  }

  fn visit_boolean_expression(&mut self, boolean_expression: &ASTBooleanExpression) {
    self.print_with_indent(&format!("Boolean: {:#?}", boolean_expression.value));
  }
}

impl ASTPrinter {
  fn print_with_indent(&mut self, text: &str) {
    println!("{}{}", " ".repeat(self.indent), text);
  }
}

pub enum ASTStatementKind {
  Expression(ASTExpression),
}

pub struct ASTStatement {
  pub kind: ASTStatementKind,
}

impl ASTStatement {
  pub fn new(kind: ASTStatementKind) -> Self {
    Self { kind }
  }

  pub fn expression(expr: ASTExpression) -> Self {
    Self::new(ASTStatementKind::Expression(expr))
  }
}

pub enum ASTExpressionKind {
  Number(ASTNumberExpression),               // 34
  Binary(ASTBinaryExpression),               // 4 + 2 * 3 + 1
  Parenthesized(ASTParenthesizedExpression), // (4 + 2) * 3 + 1
  Unary(ASTUnaryExpression),                 // - 负号，! 非 一元运算符
  Boolean(ASTBooleanExpression),             // 布尔值 true/false
}

pub struct ASTBooleanExpression {
  pub value: bool,
  pub token: Token,
}

#[derive(Debug)]
pub enum ASTUnaryOperatorKind {
  Minus,      // - 负号运算符
  BitwiseNot, // ! 非运算符
}

pub struct ASTUnaryOperator {
  pub kind: ASTUnaryOperatorKind,
  pub token: Token,
}

impl ASTUnaryOperator {
  pub fn new(kind: ASTUnaryOperatorKind, token: Token) -> Self {
    Self { kind, token }
  }
}

pub struct ASTUnaryExpression {
  pub operator: ASTUnaryOperator,  // 负号运算符
  pub operand: Box<ASTExpression>, // 操作数
}

#[derive(Debug)]
pub enum ASTBinaryOperatorKind {
  Plus,     // 加法
  Minus,    // 减法
  Multiply, // 乘法
  Divide,   // 除法
}

pub struct ASTBinaryOperator {
  pub kind: ASTBinaryOperatorKind,
  token: Token,
}

impl ASTBinaryOperator {
  pub fn new(kind: ASTBinaryOperatorKind, token: Token) -> Self {
    Self { kind, token }
  }

  pub fn precedence(&self) -> u8 {
    match self.kind {
      ASTBinaryOperatorKind::Plus | ASTBinaryOperatorKind::Minus => 1,
      ASTBinaryOperatorKind::Multiply | ASTBinaryOperatorKind::Divide => 2,
    }
  }
}

pub struct ASTNumberExpression {
  pub number: f64,
}

pub struct ASTBinaryExpression {
  pub lhs: Box<ASTExpression>,
  pub operator: ASTBinaryOperator,
  pub rhs: Box<ASTExpression>,
}

pub struct ASTParenthesizedExpression {
  expression: Box<ASTExpression>,
}

pub struct ASTExpression {
  pub kind: ASTExpressionKind,
}

impl ASTExpression {
  pub fn new(kind: ASTExpressionKind) -> Self {
    Self { kind }
  }

  // 34
  pub fn number(number: f64) -> Self {
    Self::new(ASTExpressionKind::Number(ASTNumberExpression { number }))
  }

  // 4 + 2 * 3
  pub fn binary(operator: ASTBinaryOperator, lhs: Self, rhs: Self) -> Self {
    Self::new(ASTExpressionKind::Binary(ASTBinaryExpression {
      lhs: Box::new(lhs),
      operator,
      rhs: Box::new(rhs),
    }))
  }

  // (4 + 2) * 3
  pub fn parenthesized(expression: Self) -> Self {
    Self::new(ASTExpressionKind::Parenthesized(
      ASTParenthesizedExpression {
        expression: Box::new(expression),
      },
    ))
  }

  // -  !
  pub fn unary(operator: ASTUnaryOperator, operand: ASTExpression) -> Self {
    ASTExpression::new(ASTExpressionKind::Unary(ASTUnaryExpression {
      operator,
      operand: Box::new(operand),
    }))
  }

  // true false
  pub fn boolean(token: Token, value: bool) -> Self {
    ASTExpression::new(ASTExpressionKind::Boolean(ASTBooleanExpression {
      token,
      value,
    }))
  }
}

// ------------------ AST end --------------------

pub struct Parser {
  tokens: Vec<Token>,
  current: usize,
}

impl Parser {
  pub fn new(tokens: Vec<Token>) -> Self {
    Self { tokens, current: 0 }
  }

  pub fn print_ast(&mut self) -> Ast {
    let mut ast = Ast::new();

    while let Some(stmt) = self.next_statement() {
      ast.add_statement(stmt);
    }

    ast.visualize();

    ast
  }

  pub fn next_statement(&mut self) -> Option<ASTStatement> {
    return self.parse_statement();
  }

  fn parse_statement(&mut self) -> Option<ASTStatement> {
    let token = self.current_token()?;
    if token.kind == TokenKind::Eof {
      return None;
    }

    let expr = self.parse_expression()?;

    return Some(ASTStatement::expression(expr));
  }

  fn parse_expression(&mut self) -> Option<ASTExpression> {
    // let token = self.consume_token()?;

    // match token.kind {
    //   TokenKind::Number(number) => {
    //     // token.span.literal.parse::<f64>().unwrap();
    //     Some(ASTExpression::number(number))
    //   }
    //   _ => return None,
    // }
    return self.parse_binary_expression(0);
  }

  fn parse_binary_expression(&mut self, precedence: u8) -> Option<ASTExpression> {
    // let mut lhs = self.parse_primary_expression()?;
    let mut lhs = self.parse_unary_expression()?;

    while let Some(operator) = self.parse_binary_operator() {
      let operator_precedence = operator.precedence();
      if operator_precedence < precedence {
        break;
      }

      self.consume_token()?;
      let rhs = self.parse_binary_expression(operator_precedence)?;
      lhs = ASTExpression::binary(operator, lhs, rhs);
    }

    return Some(lhs);
  }

  fn parse_unary_expression(&mut self) -> Option<ASTExpression> {
    if let Some(operator) = self.parse_unary_opreator() {
      self.consume_token()?;
      // 获取到操作数
      let opreand = self.parse_unary_expression().unwrap();
      // 传递负号和操作数
      return Some(ASTExpression::unary(operator, opreand));
    }
    return self.parse_primary_expression();
  }

  fn parse_unary_opreator(&mut self) -> Option<ASTUnaryOperator> {
    let token = self.current_token()?;

    let kind = match token.kind {
      TokenKind::Minus => Some(ASTUnaryOperatorKind::Minus), // 负号 -
      TokenKind::Bang => Some(ASTUnaryOperatorKind::BitwiseNot), // 非 !

      _ => None,
    };

    return kind.map(|kind| ASTUnaryOperator::new(kind, token.clone()));
  }

  fn parse_binary_operator(&mut self) -> Option<ASTBinaryOperator> {
    let token = self.current_token()?;
    let kind = match token.kind {
      TokenKind::Plus => Some(ASTBinaryOperatorKind::Plus),
      TokenKind::Minus => Some(ASTBinaryOperatorKind::Minus),
      TokenKind::Asterisk => Some(ASTBinaryOperatorKind::Multiply),
      TokenKind::Slash => Some(ASTBinaryOperatorKind::Divide),
      _ => None,
    };

    return kind.map(|kind| ASTBinaryOperator::new(kind, token.clone()));
  }

  fn parse_primary_expression(&mut self) -> Option<ASTExpression> {
    let token = self.consume_token()?;

    match token.kind {
      TokenKind::Number(number) => {
        // token.span.literal.parse::<f64>().unwrap();
        Some(ASTExpression::number(number))
      }
      TokenKind::LeftParen => {
        let expr = self.parse_expression()?;
        let token = self.consume_token()?;
        if token.kind != TokenKind::RightParen {
          panic!("Expected right paren");
        }
        Some(ASTExpression::parenthesized(expr))
      }
      TokenKind::True | TokenKind::False => {
        let value = token.kind == TokenKind::True;
        Some(ASTExpression::boolean(token.clone(), value))
      }
      _ => return None,
    }
  }

  fn peek_token(&self, offset: isize) -> Option<&Token> {
    let mut index = (self.current as isize + offset) as usize;
    if index >= self.tokens.len() {
      index = self.tokens.len() - 1;
    }
    self.tokens.get(index)
  }

  fn current_token(&self) -> Option<&Token> {
    self.peek_token(0)
  }

  fn consume_token(&mut self) -> Option<&Token> {
    self.current += 1;
    self.peek_token(-1)
  }
}
